from FlagEmbedding import FlagModel
import os
from my_common import BGE_LARGE_ZH_V1_5_MODEL

# 如果需要指定 GPU，可以设置环境变量
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 使用第一个 GPU

corpus = [
    "Michael Jackson was a legendary pop icon known for his record-breaking music and dance innovations.",
    "Fei-Fei Li is a professor in Stanford University, revolutionized computer vision with the ImageNet project.",
    "Brad Pitt is a versatile actor and producer known for his roles in films like 'Fight Club' and 'Once Upon a Time in Hollywood.'",
    "Geoffrey Hinton, as a foundational figure in AI, received Turing Award for his contribution in deep learning.",
    "Eminem is a renowned rapper and one of the best-selling music artists of all time.",
    "Taylor Swift is a Grammy-winning singer-songwriter known for her narrative-driven music.",
    "Sam Altman leads OpenAI as its CEO, with astonishing works of GPT series and pursuing safe and beneficial AI.",
    "Morgan Freeman is an acclaimed actor famous for his distinctive voice and diverse roles.",
    "Andrew Ng spread AI knowledge globally via public courses on Coursera and Stanford University.",
    "Robert Downey Jr. is an iconic actor best known for playing Iron Man in the Marvel Cinematic Universe.",
]
queries = ["Who could be an expert of neural network?"]

if __name__ == '__main__':
    # get the BGE embedding model
    model = FlagModel(BGE_LARGE_ZH_V1_5_MODEL,
                  query_instruction_for_retrieval="Represent this sentence for searching relevant passages:",
                  use_fp16=True)

    # 获取查询嵌入
    corpus_embeddings = model.encode(queries)

    # 获取文档嵌入
    query_embedding = model.encode(corpus)

    # 计算相似度. 计算query_embedding 和corpus_embeddings 中每个句子的相似度
    sim_scores = query_embedding @ corpus_embeddings.T

    print("查询嵌入形状:", corpus_embeddings.shape)
    print("文档嵌入形状:", query_embedding.shape)
    print("相似度分数:", sim_scores)

    print("query_embedding向量前10个元素:")
    print(query_embedding[:10])

    # get the indices in sorted order
    sorted_indices = sorted(range(len(sim_scores)), key=lambda k: sim_scores[k], reverse=True)
    print(sorted_indices)

    # iteratively print the score and corresponding sentences in descending order
    for i in sorted_indices:
        print("Score of ")
        print(sim_scores[i])
        print(corpus[i])

# Step 4: Evaluate
def evaluate():
    queries = [
        "Who could be an expert of neural network?",
        "Who might had won Grammy?",
        "Won Academy Awards",
        "One of the most famous female singers.",
        "Inventor of AlexNet",
    ]
